# -*- coding: utf-8 -*-
# ScikitLearnBase.py
"""
Scikit-learn实验基类模块
提供实验常用的通用功能，包括数据加载、预处理、可视化、结果保存等
"""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import seaborn as sns

# 设置中文字体支持 - 更健壮的字体设置
def setup_chinese_font():
    """设置中文字体支持，尝试多种字体"""
    # 尝试多种可能的中文字体
    chinese_fonts = [
        'Microsoft YaHei',  # 微软雅黑
        'SimHei',           # 黑体
        'SimSun',           # 宋体
        'KaiTi',            # 楷体
        'FangSong',         # 仿宋
        'Arial Unicode MS', # Mac上的中文字体
        'DejaVu Sans',      # 跨平台字体
        'sans-serif'        # 最后回退到无衬线字体
    ]
    
    # 设置字体
    plt.rcParams['font.sans-serif'] = chinese_fonts
    plt.rcParams['axes.unicode_minus'] = False
    
    # 设置seaborn样式
    sns.set_style("whitegrid")
    sns.set(font=chinese_fonts[0])  # 使用第一个中文字体

# 初始化字体设置
setup_chinese_font()

class ScikitLearnExperiment:
    """Scikit-learn实验基类"""
    
    def __init__(self, experiment_name, save_dir="Scikit-Learning/img"):
        """
        初始化实验
        
        参数:
        experiment_name: 实验名称
        save_dir: 图片保存目录
        """
        self.experiment_name = experiment_name
        self.save_dir = save_dir
        self.X = None
        self.y = None
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
        self.scaler = StandardScaler()
        self.feature_names = None
        self.target_names = None
        
        # 创建保存目录
        os.makedirs(self.save_dir, exist_ok=True)
        
        # 确保字体设置正确
        setup_chinese_font()
        
    def print_header(self):
        """打印实验标题"""
        print("=" * 60)
        print(f"实验：{self.experiment_name}")
        print("=" * 60)
        
    def print_footer(self):
        """打印实验结束信息"""
        print("=" * 60)
        print(f"实验顺利完成!")
        print("=" * 60)
        
    def load_dataset(self, dataset_loader, **kwargs):
        """
        加载数据集
        
        参数:
        dataset_loader: 数据集加载函数，如load_iris, load_diabetes等
        **kwargs: 传递给加载函数的参数
        """
        print(f"正在加载数据集...")
        dataset = dataset_loader(**kwargs)
        self.X = dataset.data
        self.y = dataset.target
        self.feature_names = dataset.feature_names
        self.target_names = dataset.target_names if hasattr(dataset, 'target_names') else None
        
        print(f"数据集形状: {self.X.shape}")
        if self.target_names is not None:
            print(f"目标类别: {self.target_names}")
        print(f"特征数量: {len(self.feature_names)}")
        
        return dataset
    
    def load_csv_dataset(self, file_path, target_column=None, feature_columns=None):
        """
        从CSV文件加载数据集
        
        参数:
        file_path: CSV文件路径
        target_column: 目标列名
        feature_columns: 特征列名列表，如果为None则使用所有列（除了目标列）
        """
        print(f"正在加载CSV数据集: {file_path}")
        try:
            df = pd.read_csv(file_path)
            print(f"成功加载数据集，形状: {df.shape}")
            
            if target_column is not None:
                self.y = df[target_column].values
                if feature_columns is None:
                    # 使用除了目标列外的所有列作为特征
                    feature_columns = [col for col in df.columns if col != target_column]
                self.X = df[feature_columns].values
                self.feature_names = feature_columns
            else:
                self.X = df.values
                self.feature_names = df.columns.tolist()
                
            return df
        except FileNotFoundError:
            print(f"文件未找到: {file_path}")
            raise
    
    def split_data(self, test_size=0.2, random_state=42):
        """划分训练集和测试集"""
        print(f"划分训练集和测试集 (测试集比例: {test_size})")
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            self.X, self.y, test_size=test_size, random_state=random_state
        )
        print(f"训练集样本数: {self.X_train.shape[0]}")
        print(f"测试集样本数: {self.X_test.shape[0]}")
        
    def standardize_data(self, fit_on_train=True):
        """标准化数据"""
        print("正在进行数据标准化...")
        if fit_on_train:
            self.X_train = self.scaler.fit_transform(self.X_train)
            self.X_test = self.scaler.transform(self.X_test)
        else:
            X_scaled = self.scaler.fit_transform(self.X)
            self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
                X_scaled, self.y, test_size=0.2, random_state=42
            )
        print("标准化完成!")
        
    def save_plot(self, filename, dpi=300, bbox_inches='tight'):
        """
        保存当前图表
        
        参数:
        filename: 文件名
        dpi: 分辨率
        bbox_inches: 边界框设置
        """
        # 确保字体设置正确
        setup_chinese_font()
        
        path = os.path.join(self.save_dir, filename)
        plt.savefig(path, dpi=dpi, bbox_inches=bbox_inches)
        print(f"图表已保存为 '{path}'")
        return path
        
    def plot_scatter(self, x, y, title, xlabel, ylabel, alpha=0.7, color=None, label=None):
        """
        绘制散点图
        
        参数:
        x: x轴数据
        y: y轴数据
        title: 图表标题
        xlabel: x轴标签
        ylabel: y轴标签
        alpha: 透明度
        color: 颜色
        label: 图例标签
        """
        # 确保字体设置正确
        setup_chinese_font()
        
        plt.scatter(x, y, alpha=alpha, c=color, label=label)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        if label:
            plt.legend()
        plt.grid(alpha=0.3)
        
    def plot_line(self, x, y, title, xlabel, ylabel, color='b', marker='o', linestyle='-', label=None):
        """
        绘制折线图
        
        参数:
        x: x轴数据
        y: y轴数据
        title: 图表标题
        xlabel: x轴标签
        ylabel: y轴标签
        color: 颜色
        marker: 标记样式
        linestyle: 线型
        label: 图例标签
        """
        # 确保字体设置正确
        setup_chinese_font()
        
        plt.plot(x, y, color=color, marker=marker, linestyle=linestyle, label=label)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        if label:
            plt.legend()
        plt.grid(alpha=0.3)
        
    def plot_bar(self, x, height, title, xlabel, ylabel, color='skyblue', alpha=0.7):
        """
        绘制条形图
        
        参数:
        x: x轴数据
        height: 条形高度
        title: 图表标题
        xlabel: x轴标签
        ylabel: y轴标签
        color: 颜色
        alpha: 透明度
        """
        # 确保字体设置正确
        setup_chinese_font()
        
        plt.bar(x, height, alpha=alpha, color=color)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        plt.title(title)
        plt.grid(alpha=0.3)
        
    def evaluate_model(self, model, X_test, y_test, model_name=""):
        """
        评估模型性能
        
        参数:
        model: 训练好的模型
        X_test: 测试集特征
        y_test: 测试集标签
        model_name: 模型名称
        """
        from sklearn.metrics import accuracy_score, mean_squared_error, r2_score
        
        y_pred = model.predict(X_test)
        
        # 分类问题
        if len(np.unique(y_test)) <= 20:  # 假设类别数少于20为分类问题
            accuracy = accuracy_score(y_test, y_pred)
            print(f"{model_name} 准确率: {accuracy:.4f}")
            return accuracy
        # 回归问题
        else:
            mse = mean_squared_error(y_test, y_pred)
            r2 = r2_score(y_test, y_pred)
            print(f"{model_name} 均方误差(MSE): {mse:.4f}")
            print(f"{model_name} 决定系数(R²): {r2:.4f}")
            return mse, r2
            
    def cross_validate(self, model, X, y, cv=5, scoring=None):
        """
        交叉验证评估模型
        
        参数:
        model: 模型
        X: 特征数据
        y: 目标数据
        cv: 交叉验证折数
        scoring: 评分标准
        """
        from sklearn.model_selection import cross_val_score
        
        scores = cross_val_score(model, X, y, cv=cv, scoring=scoring)
        print(f"交叉验证得分 ({scoring}): {scores.mean():.4f} (±{scores.std():.4f})")
        return scores
        
    def create_feature_name_mapping(self, mapping_dict):
        """
        创建特征名称映射（中英文对照）
        
        参数:
        mapping_dict: 映射字典，格式为{英文名: 中文名}
        """
        return [mapping_dict.get(name, name) for name in self.feature_names]
        
    def create_target_name_mapping(self, mapping_dict):
        """
        创建目标名称映射（中英文对照）
        
        参数:
        mapping_dict: 映射字典，格式为{英文名: 中文名}
        """
        if self.target_names is not None:
            return [mapping_dict.get(name, name) for name in self.target_names]
        return None